import numpy as np
import streamlit as st
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import pinecone
from PIL import Image
import os

# 初始化Pinecone（已填入API密钥，请补充环境参数）
PINECONE_API_KEY = "pcsk_TCZLb_n7yFoy1mV8YzXEuw4kTbQQ1jp9vAuKozVAP5C71xgekAKUd4fLAMSBaxEDaRoi"
PINECONE_ENV = "us-east-1"  
INDEX_NAME = "knn-digit-recognition"

# 初始化Pinecone客户端
pinecone.init(api_key=PINECONE_API_KEY, environment=PINECONE_ENV)

# 1. 准备数据并上传到Pinecone（首次运行执行）
def prepare_and_upload_data():
    # 加载数据集
    data = load_digits()
    X, y = data.data, data.target
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
    
    # 检查索引是否存在，不存在则创建
    if INDEX_NAME not in pinecone.list_indexes():
        pinecone.create_index(INDEX_NAME, dimension=X_train.shape[1])
    
    # 连接到索引
    index = pinecone.Index(INDEX_NAME)
    
    # 准备要上传的数据（格式：(id, vector, metadata)）
    vectors = []
    for i, (vector, label) in enumerate(zip(X_train, y_train)):
        vectors.append((
            f"train_{i}",  # 唯一ID
            vector.tolist(),  # 向量数据
            {"label": int(label)}  # 元数据（存储标签）
        ))
    
    # 批量上传数据（每次上传100条）
    batch_size = 100
    for i in range(0, len(vectors), batch_size):
        batch = vectors[i:i+batch_size]
        index.upsert(vectors=batch)
    
    print(f"已上传 {len(vectors)} 条训练数据到Pinecone")
    print(f"测试集准确率评估中...")
    
    # 简单评估测试集准确率
    correct = 0
    for vector, true_label in zip(X_test, y_test):
        result = index.query(vector=vector.tolist(), top_k=5, include_metadata=True)
        # 投票决定预测结果
        labels = [match["metadata"]["label"] for match in result["matches"]]
        pred_label = max(set(labels), key=labels.count)
        if pred_label == true_label:
            correct += 1
    
    accuracy = correct / len(X_test)
    print(f"测试集准确率: {accuracy:.4f}")

# 2. WebAPP界面与预测逻辑
def main():
    # 首次运行时取消注释以准备并上传数据，之后可注释
    # prepare_and_upload_data()
    
    # 连接到Pinecone索引
    if INDEX_NAME not in pinecone.list_indexes():
        st.error("索引不存在，请先运行prepare_and_upload_data()函数")
        return
    index = pinecone.Index(INDEX_NAME)
    
    st.title("手写数字识别（Pinecone KNN版）")
    
    # 手写画布
    canvas_result = st_canvas(
        fill_color="black", stroke_color="white", stroke_width=20,
        width=280, height=280, drawing_mode="freedraw", key="canvas"
    )
    
    if canvas_result.image_data is not None:
        # 图像预处理 - 转换为与训练数据一致的格式
        img = Image.fromarray(canvas_result.image_data.astype(np.uint8)).convert('L')
        img = img.resize((8, 8))  # 适配load_digits的8×8尺寸
        img_array = np.array(img).flatten() / 255.0  # 归一化到0-1
        # 转换为与训练数据相同的尺度（load_digits的像素值范围是0-16）
        img_array = img_array * 16
        
        # 使用Pinecone进行KNN搜索（取5个最近邻）
        top_k = 5
        result = index.query(
            vector=img_array.tolist(),
            top_k=top_k,
            include_metadata=True
        )
        
        # 从结果中提取标签并进行投票
        if result["matches"]:
            labels = [match["metadata"]["label"] for match in result["matches"]]
            distances = [match["score"] for match in result["matches"]]
            
            # 计算置信度（使用距离加权）
            label_counts = {}
            for label, dist in zip(labels, distances):
                # 距离越近权重越高（这里用1/(距离+1e-6)避免除零）
                weight = 1 / (dist + 1e-6)
                if label in label_counts:
                    label_counts[label] += weight
                else:
                    label_counts[label] = weight
            
            # 确定最终预测结果
            prediction = max(label_counts, key=label_counts.get)
            total_weight = sum(label_counts.values())
            confidence = label_counts[prediction] / total_weight
            
            st.write(f"预测结果：{prediction}")
            st.write(f"置信度：{confidence:.2f}")
            
            # 显示最近邻信息（可选）
            with st.expander("查看最近邻匹配详情"):
                for i, (match, label) in enumerate(zip(result["matches"], labels)):
                    st.write(f"第{i+1}近邻: 标签={label}, 距离={match['score']:.4f}")

if __name__ == "__main__":
    main()